import pyaudio
import webrtcvad
import numpy as np
from pypinyin import pinyin, Style
import re
import time
from t import ASR_API

# 用于匹配关键词的拼音
def extract_chinese_and_convert_to_pinyin(input_string):
    """
    提取字符串中的汉字，并将其转换为拼音。
    
    :param input_string: 原始字符串
    :return: 转换后的拼音字符串
    """
    # 使用正则表达式提取所有汉字
    chinese_characters = re.findall(r'[\u4e00-\u9fa5]', input_string)
    # 将汉字列表合并为字符串
    chinese_text = ''.join(chinese_characters)
    
    # 转换为拼音
    pinyin_result = pinyin(chinese_text, style=Style.NORMAL)
    # 将拼音列表拼接为字符串
    pinyin_text = ' '.join([item[0] for item in pinyin_result])
    
    return pinyin_text

def welcome_message():
    """
    打印一个带有颜色和闪烁效果的 ASCII 艺术图案，模拟可爱的小猫形象。
    表示成功进入对话系统
    """
    # 设置颜色和样式
    cyan = "\033[1;96m"      # 青色
    white = "\033[1;97m"     # 白色
    blink = "\033[5m"        # 闪烁效果
    reset = "\033[0m"        # 重置样式

    # 打印带有颜色和边框的欢迎信息
    print(f"\n{cyan}    ------------------------------{reset}\n")
    print(f"{cyan}{blink}")  # 启用闪烁效果

    # 打印 ASCII 艺术图案
    print("              /\\_____/\\")
    print("             /  o   o  \\")
    print("            ( ==  ^  == )")
    print("             )----O----(")
    print("            (           )")
    print("           ( (  )   (  ) )")
    print("          (__(__)___(__)__) \n")

    print(reset)  # 重置样式
    print(f"{cyan}    ------------------------------{reset}\n")
    print(f"{white}    喵喵喵!!{reset}\n")



# 参数配置
AUDIO_RATE = 16000       # 采样率：16000（支持8000, 16000, 32000或48000）
CHUNK_SIZE = 480         # 每块大小（30ms，保证为10/20/30ms的倍数）
VAD_MODE = 1             # VAD模式（0-3，数值越小越保守）

# 初始化VAD
vad = webrtcvad.Vad(VAD_MODE)


# 初始化语音识别模型
"""
try

pip install modelscope[framework]

or

pip install modelscope[dataset]
"""
# from modelscope.pipelines import pipeline
# from modelscope.utils.constant import Tasks
# inference_pipeline = pipeline(
#     task=Tasks.auto_speech_recognition,
#     model=r'D:\Downloads\SenseVoiceSmall',
#     model_revision="master",
#     device="cuda:0",
#     disable_update=True,
#     use_itn=True, # 加标点符号
# )


from funasr import AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess

sound_rec_model = AutoModel(
    model=r"D:\Downloads\SenseVoiceSmall",
    trust_remote_code=True,
    remote_code="./model.py",    
    vad_model="fsmn-vad",
    vad_kwargs={"max_single_segment_time": 30000},
    device="cuda:0",
    use_itn=True,
    disable_update=True
)
# 如果需要删除模型输出的emoji，可以参考一下链接：https://github.com/FunAudioLLM/SenseVoice/issues/113
# 如果要设置热词，可以参考：https://github.com/pengzhendong/streaming-sensevoice/blob/master/main.py


# 初始化说话人验证模型
from modelscope.pipelines import pipeline
sv_pipeline = pipeline(
    task='speaker-verification',
    model=r'D:\Downloads\speech_campplus_sv_zh-cn_3dspeaker_16k'
)

class SpeechDetector:
    def __init__(self, calibration_seconds=2, chunk_duration_ms=30):
        """
        calibration_seconds: 校准背景噪音所需时间（秒）
        chunk_duration_ms: 每块时长（毫秒）
        """
        self.calibration_seconds = calibration_seconds
        self.chunk_duration_ms = chunk_duration_ms
        self.calibrated = False
        self.amplitude_threshold = None

        # 语音缓冲区
        self.speech_buffer = bytearray()

        # 连续帧判决参数（降低短时噪音误判）
        self.speech_state = False          # 当前状态：True为语音，False为无语音
        self.consecutive_speech = 0        # 连续语音帧计数
        self.consecutive_silence = 0       # 连续静音帧计数
        self.required_speech_frames = 2    # 连续3帧语音后确认进入语音状态（约90ms）1帧为30ms
        self.required_silence_frames = 15  # 750ms 连续25帧静音后退出语音状态
        self.keyword_string = "ni hao xiao bai"           # 关键词拼音

    def calibrate(self, stream):
        """
        校准背景噪音：录制固定时长音频，计算平均幅值与标准差，从而设置自适应阈值
        """
        print("开始校准背景噪音，请保持安静...")
        amplitudes = []
        num_frames = int(self.calibration_seconds * (1000 / self.chunk_duration_ms))
        for _ in range(num_frames):
            audio_chunk = stream.read(CHUNK_SIZE, exception_on_overflow=False)
            audio_data = np.frombuffer(audio_chunk, dtype=np.int16)
            amplitudes.append(np.abs(audio_data).mean())
        mean_noise = np.mean(amplitudes)
        std_noise = np.std(amplitudes)
        # 阈值设置为均值加2倍标准差
        self.amplitude_threshold = mean_noise + 2 * std_noise
        print(f"校准完成：噪音均值={mean_noise:.2f}，标准差={std_noise:.2f}，设置阈值={self.amplitude_threshold:.2f}")
        self.calibrated = True

    def analyze_spectrum(self, audio_chunk):
        """
        通过频谱分析检测语音特性：
        1. 对音频块应用汉宁窗后计算 FFT
        2. 统计局部峰值数量（峰值必须超过均值的1.5倍）
        3. 当峰值数量大于等于设定阈值时，认为该块具有语音特征
        """
        audio_data = np.frombuffer(audio_chunk, dtype=np.int16)
        if len(audio_data) == 0:
            return False

        # 使用汉宁窗减少FFT泄露
        window = np.hanning(len(audio_data))
        windowed_data = audio_data * window

        # FFT计算得到频谱（只需正频率部分）
        spectrum = np.abs(np.fft.rfft(windowed_data))
        spectral_mean = np.mean(spectrum)

        # 统计超过均值1.5倍的局部峰值数量
        peak_count = 0
        for i in range(1, len(spectrum) - 1):
            if spectrum[i] > spectrum[i - 1] and spectrum[i] > spectrum[i + 1] and spectrum[i] > spectral_mean * 1.5:
                peak_count += 1

        return peak_count >= 3

    def is_speech(self, audio_chunk):
        """
        综合判断：先通过能量预处理（阈值）过滤低能量数据，再利用VAD和频谱分析判断语音。
        两者结合能有效降低噪音导致的误判。
        """
        # self.amplitude_threshold = 10000 # 临时设置阈值

        amplitude_threshold = self.amplitude_threshold if self.amplitude_threshold is not None else 500
        audio_data = np.frombuffer(audio_chunk, dtype=np.int16)
        amplitude = np.abs(audio_data).mean()
        if amplitude < amplitude_threshold:
            return False

        # VAD检测与频谱检测
        vad_result = vad.is_speech(audio_chunk, AUDIO_RATE)
        spectral_result = self.analyze_spectrum(audio_chunk)

        return vad_result and spectral_result

    def process_chunk(self, audio_chunk):
        """
        对每一块数据进行处理：综合能量检测、VAD、频谱分析，并采用连续帧策略实现状态平滑。
        修改后的逻辑如下：
          - 当检测到语音时（is_speech_chunk 为 True）：
                如果当前未进入语音状态且连续语音帧达到设定阈值，则将状态置为True，并打印提示；
                如果处于语音状态，则不断将当前的 audio_chunk 追加到缓冲区。
          - 当检测为非语音时：
                累计静音帧，若连续静音帧达到设定阈值，则认为语音结束，调用推理模型，
                并清空缓冲区后恢复初始状态。
        """
        is_speech_chunk = self.is_speech(audio_chunk)

        if is_speech_chunk:
            self.consecutive_speech += 1
            self.consecutive_silence = 0
            # 当未进入语音状态且连续语音帧达到阈值时，转为语音状态
            if not self.speech_state and self.consecutive_speech >= self.required_speech_frames:
                self.speech_state = True
                print("Detected Speech")
            # 如果处于语音状态，则持续追加当前块
            if self.speech_state:
                self.speech_buffer.extend(audio_chunk)
        else:
            self.consecutive_silence += 1
            self.consecutive_speech = 0
            # 当处于语音状态并且连续静音帧达到阈值时，结束语音录入
            if self.speech_state and self.consecutive_silence >= self.required_silence_frames:
                
                self.speech_buffer.extend(audio_chunk) # 包含尾部的静音
                
                self.speech_state = False # 结束语音录入

                print("正在识别...", len(self.speech_buffer))

                # 将缓冲区的数据转换为 NumPy 数组
                audio_data = bytes(self.speech_buffer)
                chinese_string = self.sound2text(audio_data)
                print("识别结果:", chinese_string)

                # 中文字符串转换为拼音
                chinese_pingyin = extract_chinese_and_convert_to_pinyin(chinese_string)
                print("拼音结果:", chinese_pingyin)


                if self.keyword_string in chinese_pingyin:
                    print("================================识别到关键词，开始识别===================================")

                    # 调用模型进行声纹识别
                    recongnize_result = sv_pipeline(["person_recording_1.wav", np.frombuffer(audio_data, dtype=np.int16)])
                    print("声纹识别结果:", recongnize_result)
                    if recongnize_result["score"] > 0.6:
                        print("===================================声纹识别成功，开始对话===================================")
                        
                        # 回复对话人
                        welcome_message()
                        time.sleep(5)

                        # 监听并识别声音
                        # 把识别到的文字喂给对话模型
                        # 输出对话结果
                        # 继续监听并识别声音
                        # 2秒钟后无声音，结束对话，回到监听关键词状态
                        def damoxing(string):
                            return f"大模型回复：{string}"

                        for recognized_text in ASR_API():
                            if recognized_text is not None:
                                damoxing(recognized_text)
                            break




                    else:
                        print("声纹识别失败，请重新说话")

                # recongnize_result = sv_pipeline(["boss.wav", np.frombuffer(audio_data, dtype=np.int16)])
                # print("声纹识别结果:", recongnize_result)


                # 清空缓冲区准备下一段语音
                self.speech_buffer = bytearray()
                print("No speech")
    
    def sound2text(self, audio_data):
        """
        直接输入音频数据进行推理
        """
        # text_string = inference_pipeline(audio_data)[0]["text"]
        # chinese_characters = re.findall(r'[\u4e00-\u9fa5]', text_string)
        # chinese_text = ''.join(chinese_characters)

        res = sound_rec_model.generate(
            input=audio_data,
            cache={},
            language="zh",  # "zh", "en", "yue", "ja", "ko", "nospeech"
            use_itn=True,
            batch_size_s=60,
            merge_vad=True,  #
            merge_length_s=15,
        )
        text_string = rich_transcription_postprocess(res[0]["text"])

        return text_string





def main():
    p = pyaudio.PyAudio()
    stream = p.open(
        format=pyaudio.paInt16,
        channels=1,
        rate=AUDIO_RATE,
        input=True,
        frames_per_buffer=CHUNK_SIZE
    )

    detector = SpeechDetector()
    # 校准背景噪音（建议在程序启动时进行）
    detector.calibrate(stream)

    print("开始监听，请开始说话...(按Ctrl+C停止)")
    try:
        while True:
            audio_chunk = stream.read(CHUNK_SIZE, exception_on_overflow=False)
            detector.process_chunk(audio_chunk)
    except KeyboardInterrupt:
        print("停止监听")
    finally:
        if detector.speech_buffer:
            detector.sound2text(bytes(detector.speech_buffer))
        stream.stop_stream()
        stream.close()
        p.terminate()

if __name__ == "__main__":
    main()
